#!/usr/bin/env python
# coding: utf-8
import sys
import os
import json
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor
import time
db = 'original'
OutputLocation = 's3://lambda-athena-log'
region_name = "cn-northwest-1"

def sql_submit(sql):
    try:
        cursor = connect(s3_staging_dir=OutputLocation,
                         region_name=region_name).cursor(schema_name=db)
        cursor.execute(sql)
        cursor.fetchall()
        corsor.close()
        return {
            'statusCode': 200,
            'body': json.dumps('success'),
            'msg': 'submit ok'
        }
    except Exception as e:
        return {
            'statusCode': 500,
            'body': json.dumps('error'),
            'msg': e
        }

def submit(sql):
    result = sql_submit(sql)
    if result["statusCode"] == 500:
        for i in range(50):
            time.sleep(1)
            result = sql_submit(sql)
    return result


def sql_getdata(sql):
    try:
        cursor = connect(s3_staging_dir=OutputLocation,
                         region_name=region_name).cursor(schema_name=db)
        data = list(cursor.execute(sql))
        cursor.fetchall()
        corsor.close()
        return data
    except Exception as e:
        print(e)



def sql_pandas_data(sql):
    try:
        cursor = connect(s3_staging_dir=OutputLocation,
                         region_name=region_name,
                         cursor_class=PandasCursor).cursor()
        data = cursor.execute(sql).as_pandas()
        cursor.close()
        return data
    except Exception as e:
        print(e)


